#include "../src/network/network_protocol.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

// Test helper macros
#define ASSERT(condition, message) \
    if (!(condition)) { \
        printf("ASSERTION FAILED: %s at %s:%d\n", message, __FILE__, __LINE__); \
        exit(1); \
    }

#define TEST_PORT 13306
#define TEST_DB_PATH "/tmp/test_network_gudb"

// Global test counters
static int tests_run = 0;
static int tests_passed = 0;

// Test helper functions
void print_test_header(const char* test_name) {
    printf("\n=== Testing %s ===\n", test_name);
    tests_run++;
}

void print_test_result(const char* test_name, bool passed) {
    if (passed) {
        printf("✓ %s PASSED\n", test_name);
        tests_passed++;
    } else {
        printf("✗ %s FAILED\n", test_name);
    }
}

// Test packet creation and manipulation
void test_packet_management() {
    print_test_header("Packet Management");
    
    // Test packet creation
    mysql_packet_t* packet = mysql_packet_create(100, 1);
    ASSERT(packet != NULL, "Packet creation failed");
    ASSERT(packet->length == 100, "Packet length incorrect");
    ASSERT(packet->sequence_id == 1, "Packet sequence ID incorrect");
    ASSERT(packet->payload != NULL, "Packet payload is NULL");
    
    // Test writing data to packet
    ASSERT(mysql_packet_write_uint8(packet, 0, 0x42) == 0, "Failed to write uint8");
    ASSERT(packet->payload[0] == 0x42, "uint8 write incorrect");
    
    ASSERT(mysql_packet_write_uint16(packet, 1, 0x1234) == 0, "Failed to write uint16");
    ASSERT(packet->payload[1] == 0x34, "uint16 write incorrect (low byte)");
    ASSERT(packet->payload[2] == 0x12, "uint16 write incorrect (high byte)");
    
    ASSERT(mysql_packet_write_uint32(packet, 3, 0x12345678) == 0, "Failed to write uint32");
    ASSERT(packet->payload[3] == 0x78, "uint32 write incorrect (byte 0)");
    ASSERT(packet->payload[4] == 0x56, "uint32 write incorrect (byte 1)");
    ASSERT(packet->payload[5] == 0x34, "uint32 write incorrect (byte 2)");
    ASSERT(packet->payload[6] == 0x12, "uint32 write incorrect (byte 3)");
    
    // Test reading data from packet
    ASSERT(mysql_packet_read_uint8(packet, 0) == 0x42, "uint8 read incorrect");
    ASSERT(mysql_packet_read_uint16(packet, 1) == 0x1234, "uint16 read incorrect");
    ASSERT(mysql_packet_read_uint32(packet, 3) == 0x12345678, "uint32 read incorrect");
    
    // Test string operations
    const char* test_string = "Hello, World!";
    int str_len = mysql_packet_write_string(packet, 10, test_string);
    ASSERT(str_len == strlen(test_string), "String write length incorrect");
    
    char* read_string = mysql_packet_read_string(packet, 10, str_len);
    ASSERT(read_string != NULL, "String read failed");
    ASSERT(strcmp(read_string, test_string) == 0, "String read content incorrect");
    free(read_string);
    
    mysql_packet_destroy(packet);
    print_test_result("Packet Management", true);
}

// Test length-encoded integer functions
void test_lenenc_integers() {
    print_test_header("Length-Encoded Integers");
    
    mysql_packet_t* packet = mysql_packet_create(100, 0);
    ASSERT(packet != NULL, "Packet creation failed");
    
    // Test small values (< 0xfb)
    size_t offset = 0;
    offset += mysql_packet_write_lenenc_int(packet, offset, 100);
    ASSERT(offset == 1, "Small lenenc int size incorrect");
    ASSERT(mysql_packet_read_lenenc_int(packet, 0) == 100, "Small lenenc int read incorrect");
    
    // Test medium values (< 0x10000)
    offset += mysql_packet_write_lenenc_int(packet, offset, 1000);
    ASSERT(offset == 4, "Medium lenenc int total size incorrect");
    ASSERT(mysql_packet_read_lenenc_int(packet, 1) == 1000, "Medium lenenc int read incorrect");
    
    // Test size calculation functions
    ASSERT(mysql_lenenc_int_size(50) == 1, "Small lenenc size calculation wrong");
    ASSERT(mysql_lenenc_int_size(1000) == 3, "Medium lenenc size calculation wrong");
    ASSERT(mysql_lenenc_int_size(100000) == 4, "Large lenenc size calculation wrong");
    ASSERT(mysql_lenenc_int_size(0x1000000ULL) == 9, "Very large lenenc size calculation wrong");
    
    mysql_packet_destroy(packet);
    print_test_result("Length-Encoded Integers", true);
}

// Test connection creation and cleanup
void test_connection_management() {
    print_test_header("Connection Management");
    
    // Create a socket pair for testing
    int sockfd[2];
    ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd) == 0, "Failed to create socket pair");
    
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    
    // Test connection creation
    mysql_connection_t* conn = mysql_connection_create(sockfd[0], &addr);
    ASSERT(conn != NULL, "Connection creation failed");
    ASSERT(conn->sockfd == sockfd[0], "Connection socket incorrect");
    ASSERT(conn->state == CONN_STATE_INIT, "Initial connection state incorrect");
    ASSERT(conn->sequence_id == 0, "Initial sequence ID incorrect");
    ASSERT(conn->authenticated == false, "Initial auth state incorrect");
    ASSERT(conn->autocommit == true, "Initial autocommit state incorrect");
    
    // Test buffer initialization
    ASSERT(mysql_connection_init_buffers(conn, 1024, 1024) == 0, "Buffer initialization failed");
    ASSERT(conn->read_buffer != NULL, "Read buffer not allocated");
    ASSERT(conn->write_buffer != NULL, "Write buffer not allocated");
    ASSERT(conn->read_buffer_size == 1024, "Read buffer size incorrect");
    ASSERT(conn->write_buffer_size == 1024, "Write buffer size incorrect");
    
    // Clean up
    mysql_connection_destroy(conn);
    close(sockfd[1]);
    
    print_test_result("Connection Management", true);
}

// Test handshake packet creation
void test_handshake_creation() {
    print_test_header("Handshake Creation");
    
    int sockfd[2];
    ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd) == 0, "Failed to create socket pair");
    
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    
    mysql_connection_t* conn = mysql_connection_create(sockfd[0], &addr);
    ASSERT(conn != NULL, "Connection creation failed");
    
    conn->connection_id = 12345;
    
    // Since we can't easily test the full handshake over sockets in a unit test,
    // we'll test the auth data generation
    mysql_generate_auth_data(conn->auth_plugin_data);
    
    // Verify auth data is generated (non-zero)
    bool has_data = false;
    for (int i = 0; i < MYSQL_AUTH_PLUGIN_DATA_LEN; i++) {
        if (conn->auth_plugin_data[i] != 0) {
            has_data = true;
            break;
        }
    }
    ASSERT(has_data, "Auth plugin data not generated");
    
    mysql_connection_destroy(conn);
    close(sockfd[1]);
    
    print_test_result("Handshake Creation", true);
}

// Test result set management
void test_result_set_management() {
    print_test_header("Result Set Management");
    
    // Create result set
    mysql_result_set_t* result = mysql_result_set_create(3);
    ASSERT(result != NULL, "Result set creation failed");
    ASSERT(result->column_count == 3, "Column count incorrect");
    ASSERT(result->fields != NULL, "Fields array not allocated");
    ASSERT(result->row_count == 0, "Initial row count incorrect");
    
    // Add field definitions
    ASSERT(mysql_result_set_add_field(result, 0, "id", MYSQL_TYPE_LONG, 10) == 0, 
           "Failed to add field 0");
    ASSERT(mysql_result_set_add_field(result, 1, "name", MYSQL_TYPE_VAR_STRING, 64) == 0, 
           "Failed to add field 1");
    ASSERT(mysql_result_set_add_field(result, 2, "score", MYSQL_TYPE_FLOAT, 8) == 0, 
           "Failed to add field 2");
    
    // Verify field definitions
    ASSERT(result->fields[0].name != NULL, "Field 0 name not set");
    ASSERT(strcmp(result->fields[0].name, "id") == 0, "Field 0 name incorrect");
    ASSERT(result->fields[0].type == MYSQL_TYPE_LONG, "Field 0 type incorrect");
    ASSERT(result->fields[0].column_length == 10, "Field 0 length incorrect");
    
    // Add rows
    char* row1_data[] = {"1", "Alice", "95.5"};
    uint64_t row1_lengths[] = {1, 5, 4};
    ASSERT(mysql_result_set_add_row(result, row1_data, row1_lengths) == 0, 
           "Failed to add row 1");
    
    char* row2_data[] = {"2", "Bob", NULL};  // Test NULL value
    uint64_t row2_lengths[] = {1, 3, 0};
    ASSERT(mysql_result_set_add_row(result, row2_data, row2_lengths) == 0, 
           "Failed to add row 2");
    
    // Verify rows
    ASSERT(result->row_count == 2, "Row count incorrect after adding rows");
    ASSERT(result->rows != NULL, "Rows array not allocated");
    ASSERT(result->rows[0] != NULL, "Row 0 not allocated");
    ASSERT(strcmp(result->rows[0][0], "1") == 0, "Row 0 col 0 incorrect");
    ASSERT(strcmp(result->rows[0][1], "Alice") == 0, "Row 0 col 1 incorrect");
    ASSERT(result->rows[1][2] == NULL, "Row 1 col 2 should be NULL");
    
    mysql_result_set_destroy(result);
    print_test_result("Result Set Management", true);
}

// Test server creation and configuration
void test_server_management() {
    print_test_header("Server Management");
    
    // Create a dummy database handle (in real tests, this would be a real GuDB instance)
    gudb* dummy_db = (gudb*)malloc(sizeof(gudb));
    ASSERT(dummy_db != NULL, "Failed to allocate dummy database");
    
    // Test server creation
    network_server_t* server = network_server_create(TEST_PORT, dummy_db);
    ASSERT(server != NULL, "Server creation failed");
    ASSERT(server->port == TEST_PORT, "Server port incorrect");
    ASSERT(server->database == dummy_db, "Server database reference incorrect");
    ASSERT(server->running == false, "Server should not be running initially");
    ASSERT(server->max_connections > 0, "Max connections not set");
    ASSERT(server->thread_count > 0, "Thread count not set");
    
    // Test configuration functions
    ASSERT(network_server_set_max_connections(server, 500) == 0, "Failed to set max connections");
    ASSERT(server->max_connections == 500, "Max connections not updated");
    
    ASSERT(network_server_set_buffer_sizes(server, 8192, 8192) == 0, "Failed to set buffer sizes");
    ASSERT(server->read_buffer_size == 8192, "Read buffer size not updated");
    ASSERT(server->write_buffer_size == 8192, "Write buffer size not updated");
    
    ASSERT(network_server_set_thread_count(server, 4) == 0, "Failed to set thread count");
    ASSERT(server->thread_count == 4, "Thread count not updated");
    
    network_server_destroy(server);
    free(dummy_db);
    
    print_test_result("Server Management", true);
}

// Test utility functions
void test_utility_functions() {
    print_test_header("Utility Functions");
    
    // Test command name function
    ASSERT(strcmp(mysql_get_command_name(COM_QUERY), "COM_QUERY") == 0, 
           "Command name for COM_QUERY incorrect");
    ASSERT(strcmp(mysql_get_command_name(COM_QUIT), "COM_QUIT") == 0, 
           "Command name for COM_QUIT incorrect");
    ASSERT(strcmp(mysql_get_command_name(COM_PING), "COM_PING") == 0, 
           "Command name for COM_PING incorrect");
    
    // Test error message function
    const char* error_msg = mysql_get_error_message(ER_ACCESS_DENIED_ERROR);
    ASSERT(error_msg != NULL, "Error message should not be NULL");
    ASSERT(strlen(error_msg) > 0, "Error message should not be empty");
    
    // Test length encoding size calculations
    ASSERT(mysql_lenenc_string_size("") == 1, "Empty string size incorrect");
    ASSERT(mysql_lenenc_string_size("test") == 5, "Short string size incorrect");  // 1 + 4
    
    print_test_result("Utility Functions", true);
}

// Test error handling
void test_error_handling() {
    print_test_header("Error Handling");
    
    // Test NULL parameter handling
    ASSERT(mysql_packet_create(0, 0) != NULL, "Should allow zero-length packets");
    
    mysql_packet_t* null_packet = NULL;
    mysql_packet_destroy(null_packet);  // Should not crash
    
    ASSERT(mysql_packet_write_uint8(NULL, 0, 0) == -1, "Should fail with NULL packet");
    ASSERT(mysql_packet_read_uint8(NULL, 0) == 0, "Should return 0 with NULL packet");
    
    // Test out-of-bounds access
    mysql_packet_t* small_packet = mysql_packet_create(5, 0);
    ASSERT(small_packet != NULL, "Small packet creation failed");
    ASSERT(mysql_packet_write_uint32(small_packet, 2, 0x12345678) == -1, 
           "Should fail with out-of-bounds write");
    mysql_packet_destroy(small_packet);
    
    // Test result set error conditions
    ASSERT(mysql_result_set_create(0) == NULL, "Should fail with zero columns");
    ASSERT(mysql_result_set_create(5000) == NULL, "Should fail with too many columns");
    
    mysql_result_set_t* result = mysql_result_set_create(2);
    ASSERT(result != NULL, "Result creation should succeed");
    ASSERT(mysql_result_set_add_field(result, 5, "test", MYSQL_TYPE_LONG, 10) == -1, 
           "Should fail with invalid field index");
    mysql_result_set_destroy(result);
    
    print_test_result("Error Handling", true);
}

// Test authentication functions
void test_authentication() {
    print_test_header("Authentication");
    
    int sockfd[2];
    ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd) == 0, "Failed to create socket pair");
    
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    
    mysql_connection_t* conn = mysql_connection_create(sockfd[0], &addr);
    ASSERT(conn != NULL, "Connection creation failed");
    
    // Test authentication
    ASSERT(mysql_authenticate_user(conn, "testuser", "authdata", "testdb") == 0, 
           "Authentication should succeed");
    ASSERT(strcmp(conn->username, "testuser") == 0, "Username not stored correctly");
    ASSERT(strcmp(conn->database, "testdb") == 0, "Database not stored correctly");
    ASSERT(conn->authenticated == true, "Authentication flag not set");
    ASSERT(conn->state == CONN_STATE_AUTHENTICATED, "Connection state not updated");
    
    mysql_connection_destroy(conn);
    close(sockfd[1]);
    
    print_test_result("Authentication", true);
}

// Integration test - basic protocol flow
void test_protocol_integration() {
    print_test_header("Protocol Integration");
    
    // This is a simplified integration test that verifies the basic protocol flow
    // In a full test, this would involve actual socket communication
    
    mysql_result_set_t* result = mysql_result_set_create(1);
    ASSERT(result != NULL, "Result set creation failed");
    
    ASSERT(mysql_result_set_add_field(result, 0, "test_col", MYSQL_TYPE_VAR_STRING, 50) == 0, 
           "Failed to add field");
    
    char* row_data[] = {"test_value"};
    uint64_t lengths[] = {10};
    ASSERT(mysql_result_set_add_row(result, row_data, lengths) == 0, 
           "Failed to add row");
    
    // Verify result set is properly formed
    ASSERT(result->column_count == 1, "Column count incorrect");
    ASSERT(result->row_count == 1, "Row count incorrect");
    ASSERT(result->fields[0].name != NULL, "Field name not set");
    ASSERT(result->rows[0][0] != NULL, "Row data not set");
    
    mysql_result_set_destroy(result);
    
    print_test_result("Protocol Integration", true);
}

// Run all tests
int main() {
    printf("GuDB Network Protocol Test Suite\n");
    printf("================================\n");
    
    // Run individual test functions
    test_packet_management();
    test_lenenc_integers();
    test_connection_management();
    test_handshake_creation();
    test_result_set_management();
    test_server_management();
    test_utility_functions();
    test_error_handling();
    test_authentication();
    test_protocol_integration();
    
    // Print final results
    printf("\n=== Test Summary ===\n");
    printf("Tests run: %d\n", tests_run);
    printf("Tests passed: %d\n", tests_passed);
    printf("Tests failed: %d\n", tests_run - tests_passed);
    
    if (tests_passed == tests_run) {
        printf("🎉 All tests PASSED!\n");
        return 0;
    } else {
        printf("❌ Some tests FAILED!\n");
        return 1;
    }
}